﻿#include "../../../includes/multimedia/rtmp/crtmpclient.h"
#include "../../../includes/QsLog/QsLog.h"
#include "../../../includes/common/common.h"
#include "../../../includes/multimedia/rtmp/rtmp.h"
#include "../../../includes/multimedia/rtmp/crtmpserver.h"

cRtmpClient::cRtmpClient(QObject *parent,
                         qintptr socketDescriptor,
                         cRtmpServer *rtmpserver)
    : QTcpSocket (parent),
      connection_mode_(RTMP_SERVER),
      stream_id_(0),
      max_chunk_size_(1024),
      peer_bandwidth_(5000000),
      acknowledgement_size_(5000000),
      m_curSocketId(socketDescriptor),
      m_rtmpserver(rtmpserver),
      connection_state_(HANDSHAKE),
      avc_sequence_header_size_(0),
      aac_sequence_header_size_(0),
      is_playing_(false),
      m_clientType(PUBLISHER),
      m_heartCount(0),
      max_gop_cache_len_(5000)
{
    if(m_curSocketId != -1)
        this->setSocketDescriptor(m_curSocketId);

    connect(this, SIGNAL(readyRead()), this, SLOT(OnReadyRead()));
    connect(this, SIGNAL(connected()), this, SLOT(onConnected()));
    connect(this, SIGNAL(disconnected()), this, SLOT(onDisconnected()));
}

cRtmpClient::~cRtmpClient()
{

}

/**
 * @brief cRtmpClient::onConnected 连接成功后的处理
 */
void cRtmpClient::onConnected()
{
    //emit sockConnected(this->socketDescriptor());
}

/**
 * @brief cRtmpClient::onDisconnected 连接断开后的处理
 */
void cRtmpClient::onDisconnected()
{
    emit sockDisConnect(m_curSocketId);
}

/**
 * @brief cRtmpClient::OnReadyRead 处理接收数据
 */
void cRtmpClient::OnReadyRead()
{
    m_dataPacket.append(this->readAll());

    m_heartCount = 0;

    // 循环解析包数据，数据中可能不只一包数据
    parsePacket();
}

qint64 cRtmpClient::send2(const char *data,qint32 size)
{
    return send(QByteArray(data,size));
}

/**
 * @brief cRtmpServer::showLog
 * @param log
 */
void cRtmpClient::showLog(QString log)
{
    if(log.isEmpty()) return;

    QLOG_INFO()<<log;
    emit sockLog(log);
}

/**
 * @brief cRtmpClient::send 封包并发送
 * @param data 要发送的数据
 *
 * @return 返回发送成功的数据
 */
qint64 cRtmpClient::send(QByteArray data)
{
    if(this->state() !=  QTcpSocket::ConnectedState ||
       data.isEmpty())
    {
        return -1;
    }

    QByteArray sendData = data;
    qint64 m_sendsize,m_totalsize;
    int m_sendcount = 0;

    m_totalsize = sendData.size();
    m_sendsize=0;

    do
    {
        QByteArray psendbytearray;

        if(sendData.size()-m_sendsize < BUF_SIZE)
            psendbytearray = sendData.mid(m_sendsize);
        else
            psendbytearray = sendData.mid(m_sendsize,BUF_SIZE);

        m_sendsize += this->write(psendbytearray);

        if(m_sendsize <= 0) m_sendcount+=1;

        // 如果5次数据都没发送成功，就跳出
        if(m_sendcount > 5)
            break;
    }
    while(m_sendsize < m_totalsize);

    if(m_sendsize > 0) m_heartCount = 0;

    return m_sendsize;
}

/**
 * @brief cRtmpClient::parsePacket 解包
 * @param packet 要解包的数据
 *
 * @return 如果数据解包成功返回真，否则返回假
 */
bool cRtmpClient::parsePacket(void)
{
    bool ret = true;

    if (m_RtmpHandshake.IsCompleted()) {
        ret = HandleChunk(m_dataPacket);
    }
    else {
        QByteArray res;
        res.resize(4096);
        int res_size = m_RtmpHandshake.Parse(m_dataPacket, res);
        if (res_size < 0) {
            ret = false;
        }

        if (res_size > 0) {
            this->send2(res.data(),res_size);
        }

        if (m_RtmpHandshake.IsCompleted()) {
            if(m_dataPacket.size() > 0) {
                ret = HandleChunk(m_dataPacket);
            }
        }
    }

    return ret;
}

bool cRtmpClient::HandleChunk(QByteArray& buffer)
{
    int ret = -1;

    do
    {
        RtmpMessage rtmp_msg;
        ret = m_RtmpChunk.Parse(buffer, rtmp_msg);
        if (ret >= 0) {
            if (rtmp_msg.IsCompleted()) {
                if (!HandleMessage(rtmp_msg)) {
                    return false;
                }
            }

            if (ret == 0) {
                break;
            }
        }
        else if (ret < 0) {
            return false;
        }
    } while (buffer.size() > 0);

    return true;
}

bool cRtmpClient::HandleMessage(RtmpMessage& rtmp_msg)
{
    bool ret = true;
    switch(rtmp_msg.type_id)
    {
        case RTMP_VIDEO:
            ret = HandleVideo(rtmp_msg);
            break;
        case RTMP_AUDIO:
            ret = HandleAudio(rtmp_msg);
            break;
        case RTMP_INVOKE:
            ret = HandleInvoke(rtmp_msg);
            break;
        case RTMP_NOTIFY:
            ret = HandleNotify(rtmp_msg);
            break;
        case RTMP_FLEX_MESSAGE:
            showLog("unsupported rtmp flex message.");
            ret = false;
            break;
        case RTMP_SET_CHUNK_SIZE:
            m_RtmpChunk.SetInChunkSize(ReadUint32BE(rtmp_msg.payload.data()));
            break;
        case RTMP_BANDWIDTH_SIZE:
            break;
        case RTMP_FLASH_VIDEO:
            showLog("unsupported rtmp flash video.");
            ret = false;
            break;
        case RTMP_ACK:
            break;
        case RTMP_ACK_SIZE:
            break;
        case RTMP_USER_EVENT:
            break;
        default:
            showLog(QString::asprintf("unkonw message type : %d", rtmp_msg.type_id));
            break;
    }

    return ret;
}

bool cRtmpClient::HandleDeleteStream()
{
    if(stream_path_ == "" || m_rtmpserver == NULL)
        return false;

    m_RtmpChunk.Clear();

    if(this->getclientType() == PUBLISHER)

    showLog(QString::asprintf("[%s] leave: %s, stream name: %s, stream path: %s",
                                     this->getclientType() == PUBLISHER ? "Publish" : "Play",
                                     app_.toStdString().c_str(),
                                     stream_name_.toStdString().c_str(),
                                     stream_path_.toStdString().c_str()));

    return m_rtmpserver->delSession(stream_path_,this);
}

bool cRtmpClient::HandleCreateStream()
{
    int stream_id = m_RtmpChunk.GetStreamId();

    AmfObjects objects;
    amf_encoder_.reset();
    amf_encoder_.encodeString("_result", 7);
    amf_encoder_.encodeNumber(amf_decoder_.getNumber());
    amf_encoder_.encodeObjects(objects);
    amf_encoder_.encodeNumber(stream_id);

    SendInvokeMessage(RTMP_CHUNK_INVOKE_ID, amf_encoder_.data());
    stream_id_ = stream_id;
    return true;
}

bool cRtmpClient::SendInvokeMessage(uint32_t csid, QByteArray payload)
{
    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_INVOKE;
    rtmp_msg.timestamp = 0;
    rtmp_msg.stream_id = stream_id_;
    rtmp_msg.payload = payload;
    rtmp_msg.length = payload.size();
    this->SendRtmpChunks(csid, rtmp_msg);

    return true;
}

void cRtmpClient::SendAcknowledgement()
{
    QByteArray data;
    data.resize(4);
    WriteUint32BE(data.data(), acknowledgement_size_);

    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_ACK_SIZE;
    rtmp_msg.payload = data;
    rtmp_msg.length = 4;
    SendRtmpChunks(RTMP_CHUNK_CONTROL_ID, rtmp_msg);
}

void cRtmpClient::SetChunkSize()
{
    m_RtmpChunk.SetOutChunkSize(max_chunk_size_);
    QByteArray data;
    data.resize(4);
    WriteUint32BE((char*)data.data(), max_chunk_size_);

    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_SET_CHUNK_SIZE;
    rtmp_msg.payload = data;
    rtmp_msg.length = 4;
    SendRtmpChunks(RTMP_CHUNK_CONTROL_ID, rtmp_msg);
}

void cRtmpClient::SetPeerBandwidth()
{
    QByteArray data;
    data.resize(5);
    WriteUint32BE(data.data(), peer_bandwidth_);
    data.data()[4] = 2;
    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_BANDWIDTH_SIZE;
    rtmp_msg.payload = data;
    rtmp_msg.length = 5;
    SendRtmpChunks(RTMP_CHUNK_CONTROL_ID, rtmp_msg);
}

bool cRtmpClient::SendNotifyMessage2(uint32_t csid, QByteArray payload,int playloadsize)
{
    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_NOTIFY;
    rtmp_msg.timestamp = 0;
    rtmp_msg.stream_id = stream_id_;
    rtmp_msg.payload = payload;
    rtmp_msg.length = playloadsize;
    SendRtmpChunks(csid, rtmp_msg);
    return true;
}

bool cRtmpClient::HandlePublish()
{
    if(m_rtmpserver == NULL)
        return false;

    AmfObjects objects;
    amf_encoder_.reset();
    amf_encoder_.encodeString("onStatus", 8);
    amf_encoder_.encodeNumber(0);
    amf_encoder_.encodeObjects(objects);

    bool is_error = false;

    if(m_rtmpserver->HasPublisher(stream_path_)) {
        is_error = true;
        objects["level"] = AmfObject("error");
        objects["code"] = AmfObject("NetStream.Publish.BadName");
        objects["description"] = AmfObject("Stream already publishing.");
    }
    else if(connection_state_ == START_PUBLISH) {
        is_error = true;
        objects["level"] = AmfObject("error");
        objects["code"] = AmfObject("NetStream.Publish.BadConnection");
        objects["description"] = AmfObject("Connection already publishing.");
    }
    /* else if(0)  {
        // 认证处理
    } */
    else {
        objects["level"] = AmfObject("status");
        objects["code"] = AmfObject("NetStream.Publish.Start");
        objects["description"] = AmfObject("Start publising.");

        m_rtmpserver->addSession(stream_path_,this);
    }

    amf_encoder_.encodeObjects(objects);
    SendInvokeMessage(RTMP_CHUNK_INVOKE_ID, amf_encoder_.data());

    if(is_error) {
        // Close ?
    }
    else {
        connection_state_ = START_PUBLISH;
    }

    m_clientType = PUBLISHER;

    m_rtmpserver->SetGopCache(stream_path_,max_gop_cache_len_);

    showLog(QString::asprintf("[%s] %d come: %s, stream name: %s, stream path: %s",
                              this->getclientType() == PUBLISHER ? "Publish" : "Play",
                              m_curSocketId,
                              app_.toStdString().c_str(),
                              stream_name_.toStdString().c_str(),
                              stream_path_.toStdString().c_str()));

    return true;
}

bool cRtmpClient::HandleConnect()
{
    if(!amf_decoder_.hasObject("app")) {
        return false;
    }

    AmfObject amfObj = amf_decoder_.getObject("app");
    app_ = amfObj.amf_string;
    if(app_ == "") {
        return false;
    }

    SendAcknowledgement();
    SetPeerBandwidth();
    SetChunkSize();

    AmfObjects objects;
    amf_encoder_.reset();
    amf_encoder_.encodeString("_result", 7);
    amf_encoder_.encodeNumber(amf_decoder_.getNumber());

    objects["fmsVer"] = AmfObject("FMS/4,5,0,297");
    objects["capabilities"] = AmfObject(255.0);
    objects["mode"] = AmfObject(1.0);
    amf_encoder_.encodeObjects(objects);
    objects.clear();
    objects["level"] = AmfObject("status");
    objects["code"] = AmfObject("NetConnection.Connect.Success");
    objects["description"] = AmfObject("Connection succeeded.");
    objects["objectEncoding"] = AmfObject(0.0);
    amf_encoder_.encodeObjects(objects);

    SendInvokeMessage(RTMP_CHUNK_INVOKE_ID, amf_encoder_.data());
    return true;
}

bool cRtmpClient::HandlePlay2()
{
    HandlePlay();
    //printf("[Play2] stream path: %s\n", stream_path_.c_str());
    return false;
}

bool cRtmpClient::HandlePlay()
{
    if(m_rtmpserver == NULL) return false;

    AmfObjects objects;
    amf_encoder_.reset();
    amf_encoder_.encodeString("onStatus", 8);
    amf_encoder_.encodeNumber(0);
    amf_encoder_.encodeObjects(objects);
    objects["level"] = AmfObject("status");
    objects["code"] = AmfObject("NetStream.Play.Reset");
    objects["description"] = AmfObject("Resetting and playing stream.");
    amf_encoder_.encodeObjects(objects);
    if(!SendInvokeMessage(RTMP_CHUNK_INVOKE_ID, amf_encoder_.data())) {
        return false;
    }

    objects.clear();
    amf_encoder_.reset();
    amf_encoder_.encodeString("onStatus", 8);
    amf_encoder_.encodeNumber(0);
    amf_encoder_.encodeObjects(objects);
    objects["level"] = AmfObject("status");
    objects["code"] = AmfObject("NetStream.Play.Start");
    objects["description"] = AmfObject("Started playing.");
    amf_encoder_.encodeObjects(objects);
    if(!SendInvokeMessage(RTMP_CHUNK_INVOKE_ID, amf_encoder_.data())) {
        return false;
    }

    amf_encoder_.reset();
    amf_encoder_.encodeString("|RtmpSampleAccess", 17);
    amf_encoder_.encodeBoolean(true);
    amf_encoder_.encodeBoolean(true);
    if(!SendNotifyMessage2(RTMP_CHUNK_DATA_ID, amf_encoder_.data(),amf_encoder_.size())) {
        return false;
    }

    connection_state_ = START_PLAY;
    m_clientType = PLAYER;

    m_rtmpserver->addSession(stream_path_,this);

    showLog(QString::asprintf("[%s] %d come: %s, stream name: %s, stream path: %s",
                                     this->getclientType() == PUBLISHER ? "Publish" : "Play",
                                     m_curSocketId,
                                     app_.toStdString().c_str(),
                                     stream_name_.toStdString().c_str(),
                                     stream_path_.toStdString().c_str()));

    return true;
}

bool cRtmpClient::SendVideoData(uint64_t timestamp, QByteArray payload, uint32_t payload_size)
{
    if (payload_size == 0) {
        return false;
    }

    RtmpMessage rtmp_msg;
    rtmp_msg.type_id = RTMP_VIDEO;
    rtmp_msg._timestamp = timestamp;
    rtmp_msg.stream_id = stream_id_;
    rtmp_msg.payload = payload;
    rtmp_msg.length = payload_size;
    this->SendRtmpChunks(RTMP_CHUNK_VIDEO_ID, rtmp_msg);

    return true;
}

bool cRtmpClient::HandleInvoke(RtmpMessage& rtmp_msg)
{
    bool ret  = true;
    amf_decoder_.reset();

    int bytes_used = amf_decoder_.decode((const char *)rtmp_msg.payload.data(), rtmp_msg.length, 1);
    if (bytes_used < 0) {
        return false;
    }

    QString method = amf_decoder_.getString();
    //showLog(QString::asprintf("[Method] %s", method.toStdString().c_str()));

    if (connection_mode_ == RTMP_SERVER) {
        if(rtmp_msg.stream_id == 0) {
            bytes_used += amf_decoder_.decode(rtmp_msg.payload.data()+bytes_used, rtmp_msg.length-bytes_used);
            if(method == "connect") {
                ret = HandleConnect();
            }
            else if(method == "createStream") {
                ret = HandleCreateStream();
            }
        }
        else if(rtmp_msg.stream_id == stream_id_) {
            bytes_used += amf_decoder_.decode((const char *)rtmp_msg.payload.data()+bytes_used, rtmp_msg.length-bytes_used, 3);
            stream_name_ = amf_decoder_.getString();
            stream_path_ = "/" + app_ + "/" + stream_name_;

            if((int)rtmp_msg.length > bytes_used) {
                bytes_used += amf_decoder_.decode((const char *)rtmp_msg.payload.data()+bytes_used, rtmp_msg.length-bytes_used);
            }

            if(method == "publish") {
                ret = HandlePublish();
            }
            else if(method == "play") {
                ret = HandlePlay();
            }
            else if(method == "play2") {
                ret = HandlePlay2();
            }
            else if(method == "DeleteStream") {
                ret = HandleDeleteStream();
            }
            else if (method == "releaseStream") {

            }
        }
    }

    return ret;
}

bool cRtmpClient::HandleVideo(RtmpMessage& rtmp_msg)
{
    if(m_rtmpserver == NULL)
        return false;

    uint8_t type = RTMP_VIDEO;
    uint8_t *payload = (uint8_t *)rtmp_msg.payload.data();
    uint32_t length = rtmp_msg.length;
    uint8_t frame_type = (payload[0] >> 4) & 0x0f;
    uint8_t codec_id = payload[0] & 0x0f;

    if(connection_mode_ == RTMP_SERVER)
    {
        if (frame_type == 1 && codec_id == RTMP_CODEC_ID_H264) {
            if (payload[1] == 0) {
                avc_sequence_header_size_ = length;
                avc_sequence_header_.resize(length);
                memcpy(avc_sequence_header_.data(), rtmp_msg.payload.data(), length);
                m_rtmpserver->SetAvcSequenceHeader(avc_sequence_header_, avc_sequence_header_size_);
                type = RTMP_AVC_SEQUENCE_HEADER;
            }
        }

        //m_rtmpserver->locksessions();
        m_rtmpserver->SendMediaData(stream_path_,type, rtmp_msg._timestamp,
                                    rtmp_msg.payload,
                                    rtmp_msg.length);
        //m_rtmpserver->unlocksessions();
    }

    return true;
}

bool cRtmpClient::HandleAudio(RtmpMessage& rtmp_msg)
{
    return true;
}

void cRtmpClient::SendRtmpChunks(uint32_t csid, RtmpMessage& rtmp_msg)
{
    uint32_t capacity = rtmp_msg.length + rtmp_msg.length/ max_chunk_size_ *5 + 1024;
    QByteArray buffer;
    buffer.resize(capacity);

    int size = m_RtmpChunk.CreateChunk(csid, rtmp_msg, buffer.data(), capacity);
    if (size > 0) {
        this->send2(buffer.data(), size);
    }
}

bool cRtmpClient::HandleNotify(RtmpMessage& rtmp_msg)
{
    amf_decoder_.reset();
    int bytes_used = amf_decoder_.decode((const char *)rtmp_msg.payload.data(), rtmp_msg.length, 1);
    if(bytes_used < 0) {
        return false;
    }

    if(amf_decoder_.getString() == "@setDataFrame")
    {
        amf_decoder_.reset();
        bytes_used = amf_decoder_.decode((const char *)rtmp_msg.payload.data()+bytes_used, rtmp_msg.length-bytes_used, 1);
        if(bytes_used < 0) {
            return false;
        }

        if(amf_decoder_.getString() == "onMetaData") {
            amf_decoder_.decode((const char *)rtmp_msg.payload.data()+bytes_used, rtmp_msg.length-bytes_used);
            meta_data_ = amf_decoder_.getObjects();

            if(m_rtmpserver)
            {
                //m_rtmpserver->locksessions();
                m_rtmpserver->SetMetaData(meta_data_);
                m_rtmpserver->SendMetaData(stream_path_,meta_data_);
                //m_rtmpserver->unlocksessions();
            }
        }
    }

    return true;
}

bool cRtmpClient::SendMetaData(AmfObjects meta_data)
{
    if (meta_data.size() == 0) {
        return false;
    }

    amf_encoder_.reset();
    amf_encoder_.encodeString("onMetaData", 10);
    amf_encoder_.encodeECMA(meta_data);

    if(!this->SendNotifyMessage2(RTMP_CHUNK_DATA_ID, amf_encoder_.data(),amf_encoder_.size())) {
        return false;
    }

    return true;
}

bool cRtmpClient::SendMediaData(uint8_t type, uint64_t timestamp,
                                QByteArray payload,int payloadsize)
{
    if (payload.isEmpty()) {
        return false;
    }

    if (type == RTMP_AVC_SEQUENCE_HEADER) {
        avc_sequence_header_ = payload;
        avc_sequence_header_size_ = payloadsize;
    }
    else if (type == RTMP_AAC_SEQUENCE_HEADER) {
        aac_sequence_header_ = payload;
        aac_sequence_header_size_ = payloadsize;
    }

    is_playing_ = true;

    RtmpMessage rtmp_msg;
    rtmp_msg._timestamp = timestamp;
    rtmp_msg.stream_id = stream_id_;
    rtmp_msg.payload = payload;
    rtmp_msg.length = payloadsize;

    if (type == RTMP_VIDEO || type == RTMP_AVC_SEQUENCE_HEADER) {
        rtmp_msg.type_id = RTMP_VIDEO;
        this->SendRtmpChunks(RTMP_CHUNK_VIDEO_ID, rtmp_msg);
    }
    else if (type == RTMP_AUDIO || type == RTMP_AAC_SEQUENCE_HEADER) {
        rtmp_msg.type_id = RTMP_AUDIO;
        this->SendRtmpChunks(RTMP_CHUNK_AUDIO_ID, rtmp_msg);
    }

    return true;
}

